"""
债券对比工作流

使用LangGraph的StateGraph实现债券对比分析流程：
- 债券对比分析
- 生成对比报告
- 使用自定义事件流与债券分析工作流保持一致
- 已移除冗余的意图识别（由主工作流负责）
"""

import logging
from typing import TypedDict, Annotated, AsyncIterator, Dict, Any, List, Optional
from datetime import datetime
import operator

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.config import get_stream_writer
from langchain_core.runnables import RunnableConfig

from agents.bond_comparison_agent import BondComparisonAgent


class ComparisonState(TypedDict):
    """债券对比状态"""
    # 输入
    bond_codes: List[str]
    user_query: str
    
    # ✅ 对比结果改为单个字符串（避免累积）
    comparison_result: str
    
    # 元数据
    errors: Annotated[List[str], operator.add]
    completed_steps: Annotated[List[str], operator.add]


class BondComparisonWorkflow:
    """债券对比工作流"""
    
    def __init__(self):
        """初始化工作流"""
        self.logger = logging.getLogger("workflows.bond_comparison")
        
        # 初始化智能体
        self.comparison_agent = BondComparisonAgent()
        
        # ✅ 移除检查点保存器，避免状态累积
        self.checkpointer = None
        
        # 构建工作流图
        self.app = self._build_workflow()
    
    def _build_workflow(self):
        """构建工作流状态图（已移除冗余的意图识别）"""
        # 创建状态图
        workflow = StateGraph(ComparisonState)
        
        # 添加节点
        workflow.add_node("comparison", self._comparison_node)
        
        # ✅ 直接从入口点执行对比（无需意图识别）
        workflow.set_entry_point("comparison")
        workflow.add_edge("comparison", END)
        
        # 编译工作流（不使用checkpointer，确保每次对比独立）
        return workflow.compile()
    
    async def _comparison_node(self, state: ComparisonState) -> Dict[str, Any]:
        """债券对比分析节点 - 使用自定义事件流"""
        try:
            bond_codes = state.get("bond_codes", [])
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 验证债券代码数量
            if len(bond_codes) < 2:
                writer({
                    "type": "error",
                    "node": "comparison",
                    "error": "请提供至少2个债券代码进行对比",
                    "timestamp": datetime.now().isoformat()
                })
                return {
                    "comparison_result": "请提供至少2个债券代码进行对比",  # ✅ 字符串
                    "errors": ["债券代码数量不足"],
                    "completed_steps": ["comparison"]
                }
            
            if len(bond_codes) > 5:
                writer({
                    "type": "error",
                    "node": "comparison",
                    "error": "最多支持对比5个债券",
                    "timestamp": datetime.now().isoformat()
                })
                return {
                    "comparison_result": "最多支持对比5个债券",  # ✅ 字符串
                    "errors": ["债券代码数量超限"],
                    "completed_steps": ["comparison"]
                }
            
            self.logger.info(f"开始对比分析: {bond_codes}")
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "comparison",
                "step_name": "📊 债券对比分析",
                "timestamp": datetime.now().isoformat()
            })
            
            # 将债券代码列表转换为逗号分隔的字符串
            bond_codes_str = ",".join(bond_codes)
            
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.comparison_agent.stream_analyze(bond_codes_str):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "comparison",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI（避免在Step中显示完整结果）
                    results.append(event.get("content", ""))
            
            # ✅ 拼接为单个字符串
            comparison_result = "\n".join(results) if results else "对比分析结果为空"
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "comparison",
                "step_name": "📊 债券对比分析",
                "timestamp": datetime.now().isoformat()
            })
            
            # ✅ 构建完整状态
            complete_state = {
                "bond_codes": bond_codes,
                "user_query": state.get("user_query", ""),
                "comparison_result": comparison_result,  # ✅ 字符串
                "errors": state.get("errors", []),
                "completed_steps": state.get("completed_steps", []) + ["comparison"]
            }
            
            # ✅ 发送final_result事件
            writer({
                "type": "final_result",
                "state": complete_state,
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "comparison_result": comparison_result,  # ✅ 字符串
                "completed_steps": ["comparison"]
            }
        except Exception as e:
            self.logger.error(f"对比分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "comparison",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "comparison_result": f"对比分析失败: {str(e)}",  # ✅ 字符串
                "errors": [f"对比分析失败: {str(e)}"],
                "completed_steps": ["comparison"]
            }
    
    async def execute(
        self, 
        bond_codes: List[str], 
        user_query: str = "",
        thread_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        执行债券对比工作流
        
        Args:
            bond_codes: 债券代码列表（2-5个）
            user_query: 用户查询（可选）
            thread_id: 线程ID（用于状态持久化）
            
        Returns:
            对比分析结果
        """
        if not thread_id:
            codes_str = "_".join(bond_codes[:3])
            thread_id = f"bond_comparison_{codes_str}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 验证债券代码数量
        if len(bond_codes) < 2:
            return {
                "status": "error",
                "error": "请提供至少2个债券代码进行对比",
                "timestamp": datetime.now().isoformat()
            }
        
        if len(bond_codes) > 5:
            return {
                "status": "error",
                "error": "最多支持对比5个债券",
                "timestamp": datetime.now().isoformat()
            }
        
        # 初始化状态
        initial_state: ComparisonState = {
            "bond_codes": bond_codes,
            "user_query": user_query or f"对比债券: {', '.join(bond_codes)}",
            "comparison_result": "",  # ✅ 空字符串
            "errors": [],
            "completed_steps": []
        }
        
        try:
            # 执行工作流
            config = RunnableConfig(configurable={"thread_id": thread_id})
            final_state = await self.app.ainvoke(initial_state, config)
            
            # 格式化结果
            return {
                "status": "success",
                "bond_codes": bond_codes,
                "timestamp": datetime.now().isoformat(),
                "results": {
                    "comparison_report": final_state.get("comparison_result", "")  # ✅ 直接使用字符串
                },
                "metadata": {
                    "completed_steps": final_state.get("completed_steps", []),
                    "errors": final_state.get("errors", [])
                }
            }
        except Exception as e:
            self.logger.error(f"工作流执行失败: {e}")
            return {
                "status": "error",
                "bond_codes": bond_codes,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    async def stream_execute(
        self,
        bond_codes: List[str],
        user_query: str = "",
        thread_id: Optional[str] = None
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        流式执行债券对比工作流（使用LangGraph自定义事件流）
        
        ✅ 与债券分析工作流保持一致：
        - 使用get_stream_writer()在节点内发送自定义事件
        - 使用stream_mode="custom"接收实时推理内容
        
        Args:
            bond_codes: 债券代码列表
            user_query: 用户查询
            thread_id: 线程ID
            
        Yields:
            工作流执行的每个步骤事件，包括智能体的实时推理内容
        """
        if not thread_id:
            codes_str = "_".join(bond_codes[:3])
            thread_id = f"bond_comparison_{codes_str}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 验证债券代码数量
        if len(bond_codes) < 2 or len(bond_codes) > 5:
            yield {
                "type": "error",
                "error": "债券代码数量应在2-5个之间",
                "timestamp": datetime.now().isoformat()
            }
            return
        
        # 初始化状态
        initial_state: ComparisonState = {
            "bond_codes": bond_codes,
            "user_query": user_query or f"对比债券: {', '.join(bond_codes)}",
            "comparison_result": "",  # ✅ 空字符串
            "errors": [],
            "completed_steps": []
        }
        
        config = RunnableConfig(configurable={"thread_id": thread_id})
        
        try:
            # ✅ 使用stream_mode="custom"接收节点内通过writer发送的自定义事件
            async for event in self.app.astream(
                initial_state,
                config,
                stream_mode="custom"  # 关键：只接收自定义事件
            ):
                # 直接转发自定义事件给UI层
                yield event
            
        except Exception as e:
            self.logger.error(f"流式执行失败: {e}", exc_info=True)
            yield {
                "type": "error",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }


# 全局实例
bond_comparison_workflow = BondComparisonWorkflow()


# 便捷函数
async def compare_bonds(bond_codes: List[str], user_query: str = "") -> Dict[str, Any]:
    """
    对比债券的便捷函数
    
    Args:
        bond_codes: 债券代码列表（2-5个）
        user_query: 用户查询
        
    Returns:
        对比结果
    """
    return await bond_comparison_workflow.execute(bond_codes, user_query)